/*-
 * Copyright (c) 2007 Microsoft
 * All rights reserved.
 */

#include "StdAfx.h"
#include "MPICore.h"
#include "Comm.h"
#include "MPIException.h"
#include "SendAsyncResult.h"
#include "TypeMapper.h"
#include "Operation.h"

using namespace System;
using namespace System::Collections;
using namespace System::Runtime::InteropServices;
using namespace Microsoft::ComputeCluster::MpiNet;

///<summary>
///Initialize a Communicators.
///</summary>
///<param name="handle">MPI Comm handle</param>
Comm::Comm(int handle)
{
    int ret, value;

    this->handle = handle;

    ret = MPI_Comm_rank(handle, &value);
    if (ret != MPI_SUCCESS)
        throw gcnew MPIException(ret);
    this->rank = value;

    ret = MPI_Comm_size(handle, &value);
    if (ret != MPI_SUCCESS)
        throw gcnew MPIException(ret);
    this->size = value;

    MPI_Errhandler_set(this->handle, MPI_ERRORS_RETURN);
}

///<summary>
///Get the channel rank
///</summary>
int Comm::Rank::get()
{
    return this->rank;
}

///<summary>
///Get the cluster size
///</summary>
int Comm::Size::get()
{
    return this->size;
}

///<summary>
///Broadcasts a value from the process with rank "root" to all other processes of the group.
///A MPI Exception will be throw if there is any MPI error.
///</summary>
///<remarks>
///It doesn't support an arry. 
///</remarks>
///<param name="root">The broadcast root</param>
///<param name="value">The value will be broadcast</param>
generic<typename T> void Comm::Bcast(int root, T% value)
{
    int ret;
    Type^ t = T::typeid;

    if (t->IsArray)
    {
        throw gcnew System::ArgumentException("Use BcastArray instead");
    }

    if (t->IsPrimitive)
    {
        T n = value;

        int mpitype = TypeMapper::Get(t);

        ret = MPI_Bcast(&n, 1, mpitype, root, handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        value = n;
        return;
    }

    if (t->IsExplicitLayout || t->IsLayoutSequential)
    {
        int length = Marshal::SizeOf(value);
        if (root == this->rank) // i am sender
        {
            GCHandle^ mem = GCHandle::Alloc(value, GCHandleType::Pinned);
            ret = MPI_Bcast((void*)mem->AddrOfPinnedObject(), length, MPI_BYTE, root, this->handle);
            mem->Free();
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);
        }
        else
        {
            void* buffer = new char[length];
            try
            {
                IntPtr^ pnt = gcnew IntPtr(buffer);
                ret = MPI_Bcast(buffer, length, MPI_BYTE, root, this->handle);
                if (ret != MPI_SUCCESS)
                    throw gcnew MPIException(ret);
                value = (T)Marshal::PtrToStructure(*pnt, t);
            }
            finally
            {
                delete[] buffer;
            }
        }
        return;
    }

    if (t->IsSerializable)
    {
        int length;
        System::Runtime::Serialization::Formatters::Binary::BinaryFormatter^ formatter = 
            gcnew System::Runtime::Serialization::Formatters::Binary::BinaryFormatter();

        if (root == this->rank) // i am sender
        {
            System::IO::MemoryStream^ memoryStream = gcnew System::IO::MemoryStream();
            formatter->Serialize(memoryStream, (Object^)value);
            length = (int)memoryStream->Length;
            ret = MPI_Bcast(&length, 1, MPI_LONG, root, this->handle);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);

            array<unsigned char>^ buffer = memoryStream->GetBuffer();
            pin_ptr<unsigned char> p = &buffer[0];
            ret = MPI_Bcast(p, length, MPI_BYTE, root, this->handle);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);
        }
        else
        {
            ret = MPI_Bcast(&length, 1, MPI_LONG, root, this->handle);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);
            unsigned char* buffer = new unsigned char[length];
            try
            {
                ret = MPI_Bcast(buffer, length, MPI_BYTE, root, this->handle);
                if (ret != MPI_SUCCESS)
                    throw gcnew MPIException(ret);
                System::IO::Stream^ memoryStream = gcnew System::IO::UnmanagedMemoryStream(buffer, length);
                value = (T)formatter->Deserialize(memoryStream);
            }
            finally
            {
                delete[] buffer;
            }
        }
        return;
    }

    throw gcnew System::NotImplementedException();
}

///<summary>
///Broadcasts a array from the process with rank "root" to all other processes of the group.
///A MPI Exception will be throw if there is any MPI error.
///</summary>
///<remarks>
///Only the array with primitive element is supported for now.
///</remarks>
///<param name="root">The broadcast root</param>
///<param name="data">The value will be broadcast</param>
///<param name="index">index of array </param>
///<param name="length">lenth</param>
generic<typename T>
void Comm::Bcast(int root, array<T>^ data, int index, int length)
{
    if (T::typeid->IsPrimitive)
    {
        pin_ptr<T> datap = &data[index];
        int mpitype = TypeMapper::Get(T::typeid);
        int ret = MPI_Bcast(datap, length, mpitype, root, handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    throw gcnew System::NotImplementedException();
}

///<summary>
///Send a array from the process with rank "root" to all other processes of the group.
///A MPI Exception will be throw if there is any MPI error.
///</summary>
///<param name="target">The send target</param>
///<param name="data">The data will be send</param>
///<param name="tag">The send tag</param>
///<param name="index">index of array </param>
///<param name="length">lenth</param>
generic<typename T>
void Comm::Send(int target, array<T>^ data, int tag, int index, int length)
{
    if (T::typeid->IsPrimitive)
    {
        pin_ptr<T> datap = &data[index];
        int mpitype = TypeMapper::Get(T::typeid);
        int ret = MPI_Send(datap, length, mpitype, target, tag, handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    throw gcnew System::NotImplementedException();
}

///<summary>
///Async Send a data from the process with rank "root" to all other processes of the group.
///</summary>
///<param name="target">The send target</param>
///<param name="data">The data will be send</param>
///<param name="tag">The send tag</param>
///<param name="callback">callback</param>
generic<typename T> 
IAsyncResult^ Comm::BeginSend(int target, T data, int tag, AsyncCallback^ callback)
{
    SendAsyncResult<T>^ r = gcnew SendAsyncResult<T>(this, target, data, tag, callback);
    return r;
}

///<summary>
///Async Send a data from the process with rank "root" to all other processes of the group.
///</summary>
///<param name="target">The send target</param>
///<param name="data">The data will be send</param>
///<param name="tag">The send tag</param>
///<param name="callback">callback</param>
generic<typename T> 
IAsyncResult^ Comm::BeginSend(int target, array<T>^ data, int tag, int index, int length, AsyncCallback^ callback)
{
    SendAsyncResult<T>^ r = gcnew SendAsyncResult<T>(this, target, data, tag, index, length, callback);
    return r;
}

///<summary>
///Send a data from the process with rank "root" to all other processes of the group.
///</summary>
///<param name="target">The send target</param>
///<param name="data">The data will be send</param>
///<param name="tag">The send tag</param>
generic<typename T>
void Comm::Send(int target, T data, int tag)
{
    Type^ t = T::typeid;
    int ret;

    if (t->IsArray)
    {
        throw gcnew System::NotSupportedException();
    }

    if (t->IsPrimitive)
    {
        int mpitype = TypeMapper::Get(t);
        ret = MPI_Send(&data, 1, mpitype, target, tag, handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    if (t->IsExplicitLayout || t->IsLayoutSequential)
    {
        int length = Marshal::SizeOf(data);
        array<char>^ buffer = gcnew array<char>(length);
        pin_ptr<char> pnt = &buffer[0];

        Marshal::StructureToPtr(data, *(gcnew IntPtr(pnt)), false);
        ret = MPI_Send(pnt, length, MPI_BYTE, target, tag, this->handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    if (t->IsSerializable)
    {
        int length;
        System::Runtime::Serialization::Formatters::Binary::BinaryFormatter^ formatter = 
            gcnew System::Runtime::Serialization::Formatters::Binary::BinaryFormatter();

        System::IO::MemoryStream^ memoryStream = gcnew System::IO::MemoryStream();
        formatter->Serialize(memoryStream, (Object^)data);
        length = (int)memoryStream->Length;
        ret = MPI_Send(&length, 1, MPI_LONG, target, tag, this->handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);

        array<unsigned char>^ buffer = memoryStream->GetBuffer();
        pin_ptr<unsigned char> p = &buffer[0];
        ret = MPI_Send(p, length, MPI_BYTE, target, tag, this->handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    throw gcnew System::NotSupportedException();
}

///<summary>
///Receive a array from target process
///</summary>
///<param name="target">target process</param>
///<param name="data">The array contains the receive data</param>
///<param name="tag">The receive tag</param>
///<param name="index">index of array </param>
///<param name="length">lenth</param>
generic<typename t>
void Comm::Receive(int target, array<t>^ data, int tag, int index, int length)
{
    if (t::typeid->IsPrimitive)
    {
        pin_ptr<t> datap = &data[index];
        int mpitype = TypeMapper::Get(t::typeid);
        MPI_Status status;
        int ret = MPI_Recv(datap, length, mpitype, target, tag, handle, &status);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    throw gcnew System::NotImplementedException();
}

///<summary>
///Receive a data from target process
///</summary>
///<param name="target">target process</param>
///<param name="data">the receive data</param>
///<param name="tag">The receive tag</param>
generic<typename T>
void Comm::Receive(int target, [Out] T% data, int tag)
{
    Type^ t = T::typeid;
    int ret;
    MPI_Status status;

    if (t->IsPrimitive)
    {
        T value;
        int mpitype = TypeMapper::Get(t);
        ret = MPI_Recv(&value, 1, mpitype, target, tag, handle, &status);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        data = value;
        return;
    }

    if (t->IsExplicitLayout || t->IsLayoutSequential)
    {
        int length = Marshal::SizeOf(data);
        void* buffer = new char[length];
        try
        {
            IntPtr^ pnt = gcnew IntPtr((char*)buffer);
            ret = MPI_Recv(buffer, length, MPI_BYTE, target, tag, this->handle, &status);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);
            data = (T)Marshal::PtrToStructure(*pnt, t);
        }
        finally{
            delete[] buffer;
        }
        return;
    }

    if (t->IsSerializable)
    {
        int length;
        System::Runtime::Serialization::IFormatter^ formatter = 
            gcnew System::Runtime::Serialization::Formatters::Binary::BinaryFormatter();

        ret = MPI_Recv(&length, 1, MPI_LONG, target, tag, this->handle, &status);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        array<unsigned char>^ buffer = gcnew array<unsigned char>(length);
        pin_ptr<unsigned char> p = &buffer[0];
        ret = MPI_Recv(p, length, MPI_BYTE, target, tag, this->handle, &status);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        System::IO::MemoryStream^ memoryStream = gcnew System::IO::MemoryStream(buffer);
        data = (T)formatter->Deserialize(memoryStream);
        return;
    }

    throw gcnew System::NotSupportedException();
}

///<summary>
///Blocks until all process have reached this routine.
///</summary>
void Comm::Barrier()
{
    int ret = MPI_Barrier(this->handle);
    if (ret != MPI_SUCCESS)
        throw gcnew MPIException(ret);
}

///<summary>
///Gathers together values from a group of processes 
///</summary>
generic<typename T> void Comm::Gather(int root, T send, array<T>^ data)
{
    Type^ t = T::typeid;
    int ret;

    if (t->IsPrimitive)
    {
        pin_ptr<T> p = &data[0];
        int mpitype = TypeMapper::Get(t);

        ret = MPI_Gather(&send, 1, mpitype, p, 1, mpitype, root, this->handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        return;
    }

    if (t->IsSerializable)
    {
        int length;
        System::Runtime::Serialization::IFormatter^ formatter = 
            gcnew System::Runtime::Serialization::Formatters::Binary::BinaryFormatter();

        if (root == this->rank)
        {
            int totalsize;
            int *recvcounts  = new int[this->Size];
            int *displs = new int[this->Size];
            array<unsigned char>^ buffer;

            try
            {
                length = 0;
                // gather the length first
                ret = MPI_Gather(&length, 1, MPI_INT, recvcounts, 1, MPI_INT, root, this->handle);
                if (ret != MPI_SUCCESS){
                    throw gcnew MPIException(ret);
                }
                recvcounts[root] = 0;
                displs[0] = 0;
                totalsize = recvcounts[this->Size - 1] + 1;
                for(int i = 1; i < this->Size; i++)
                {
                    totalsize+=recvcounts[i - 1];
                    displs[i] = displs[i - 1] + recvcounts[i - 1];
                }

                buffer = gcnew array<unsigned char>(totalsize);
                pin_ptr<unsigned char> p = &buffer[0];
                // gather the real data
                ret = MPI_Gatherv(NULL, 0, MPI_CHAR, p, recvcounts, displs, MPI_CHAR, root, this->handle);
                if (ret != MPI_SUCCESS){
                    throw gcnew MPIException(ret);
                }

                // deserialize all data
                for (int i = 0 ; i < this->Size; i++)
                {
                    if (this->Rank == i)
                    {
                        data[i] = send;
                    }
                    else
                    {
                        System::IO::MemoryStream^ memoryStream = gcnew System::IO::MemoryStream(buffer, displs[i], recvcounts[i]);
                        T tmp = (T)formatter->Deserialize(memoryStream);
                        data[i] = tmp;
                    }
                }
            }
            finally{
                delete[] recvcounts;
                delete[] displs;
            }
        }
        else
        {
            System::IO::MemoryStream^ memoryStream = gcnew System::IO::MemoryStream();
            formatter->Serialize(memoryStream, (Object^)send);
            length = (int)memoryStream->Length;

            ret = MPI_Gather(&length, 1, MPI_INT, NULL, 1, MPI_INT, root, this->handle);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);

            array<unsigned char>^ buffer = memoryStream->GetBuffer();
            pin_ptr<unsigned char> p = &buffer[0];
            ret = MPI_Gatherv(p, length, MPI_CHAR, NULL, NULL, NULL, MPI_CHAR, root, this->handle);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);

        }
        return;
    }

    throw gcnew System::NotSupportedException();
}

///<summary>
///Sends data from one task to all other tasks in a group 
///</summary>
generic<typename T> void Comm::Scatter(int root, [Out] T% recv, array<T>^ data)
{
    Type^ t = T::typeid;
    int ret;

    if (t->IsPrimitive)
    {
        T value;
        pin_ptr<T> p = &data[0];
        int mpitype = TypeMapper::Get(t);

        ret = MPI_Scatter(p, 1, mpitype, &value, 1, mpitype, root, this->handle);
        if (ret != MPI_SUCCESS)
            throw gcnew MPIException(ret);
        recv = value;
        return;
    }

    if (t->IsSerializable)
    {
        int length;
        System::Runtime::Serialization::IFormatter^ formatter = 
            gcnew System::Runtime::Serialization::Formatters::Binary::BinaryFormatter();

        if (root == this->rank)
        {
            int *sendlength = new int[this->Size];
            int *displs = new int[this->Size];
            char* buffer = NULL;

            try
            {
                array<System::IO::MemoryStream^>^ datas = gcnew array<System::IO::MemoryStream^>(this->Size);
                unsigned totalLength = 0;

                for (int i = 0; i < this->Size; i++)
                {
                    if (i != this->Rank)
                    {
                        System::IO::MemoryStream^ memoryStream = gcnew System::IO::MemoryStream();
                        formatter->Serialize(memoryStream, (Object^)data[i]);
                        sendlength[i] = (int)memoryStream->Length;
                        totalLength += sendlength[i];
                        datas[i] = memoryStream;
                    }
                    else
                    {
                        sendlength[i] = 0;
                    }
                }
                ret = MPI_Scatter(sendlength, 1, MPI_INT, &length, 1, MPI_INT, root, this->handle);
                if (ret != MPI_SUCCESS)
                {
                    throw gcnew MPIException(ret);
                }

                buffer = new char[totalLength];
                displs[0] = 0;
                for(int i = 0; i < this->Size; i++)
                {
                    if (i != this->Rank)
                    {
                        array<unsigned char>^ buf = datas[i]->GetBuffer();
                        pin_ptr<unsigned char> p = &buf[0];

                        // memcpy(p, buffer + displs[i], sendlength[i]);
                        for (int k = 0; k < sendlength[i]; k++)
                        {
                            *(buffer + displs[i] + k) = *(p + k);
                        }

                    }
                    if (i != this->Size - 1)
                        displs[i + 1] = displs[i] + sendlength[i];
                }

                ret = MPI_Scatterv(buffer, sendlength, displs, MPI_CHAR, NULL, 0, MPI_CHAR, root, this->handle);
                if (ret != MPI_SUCCESS)
                    throw gcnew MPIException(ret);

                recv = data[root];
            }
            finally{
                if (buffer)
                    delete[] buffer;

                delete[] displs;
                delete[] sendlength;
            }
            return;
        }
        else
        {
            ret = MPI_Scatter(NULL, 1, MPI_INT, &length, 1, MPI_INT, root, this->handle);
            if (ret != MPI_SUCCESS)
                throw gcnew MPIException(ret);

            unsigned char *buffer = new unsigned char[length];
            try
            {
                ret = MPI_Scatterv(NULL, NULL, NULL, MPI_CHAR, buffer, length, MPI_CHAR, root, this->handle);
                if (ret != MPI_SUCCESS)
                    throw gcnew MPIException(ret);

                System::IO::Stream^ memoryStream = gcnew System::IO::UnmanagedMemoryStream(buffer, length);
                recv = (T)formatter->Deserialize(memoryStream);
            }
            finally
            {
                delete[] buffer;
            }
        }


        return;
    }

    throw gcnew System::NotSupportedException();
}

void Comm::EndSend(IAsyncResult^ result)
{
    AsyncResult^ r = dynamic_cast<AsyncResult^>(result);

    if (r == nullptr)
        throw gcnew System::InvalidOperationException();

    r->WaitComplete();
}